/*
 *  Copyright 2008-2009 NVIDIA Corporation
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

/*! \file jacobi.inl
 *  \brief Inline file for jacobi.h
 */

#include <cusp/multiply.h>
#include <cusp/detail/format_utils.h>
#include <cusp/relaxation/jacobi.h>

#include <thrust/functional.h>
#include <thrust/transform.h>
#include <thrust/iterator/zip_iterator.h>

namespace cusp
{
namespace relaxation
{

// constructor
template <typename ValueType, typename MemorySpace>
multi_jacobi<ValueType,MemorySpace>
::multi_jacobi() : default_omega(0.0)
{
}

template <typename ValueType, typename MemorySpace>
template<typename MatrixType>
multi_jacobi<ValueType,MemorySpace>
::multi_jacobi(const MatrixType& A, ValueType omega)
    : default_omega(omega), temp(A.partition, 0), diagonal(A.partition, 0), x2(A.num_rows, 0)
{
    CUSP_PROFILE_SCOPED();

    // extract the main diagonal
    cusp::detail::extract_diagonal(A, diagonal);

}

template <typename ValueType, typename MemorySpace>
template<typename MatrixType, typename PartitionType>
multi_jacobi<ValueType,MemorySpace>
::multi_jacobi(const MatrixType& A, ValueType omega, PartitionType &partition)
    : default_omega(omega), temp(partition, 0), diagonal(partition, 0), x2(A.num_rows, 0)
{
    CUSP_PROFILE_SCOPED();

    // extract the main diagonal
    cusp::detail::extract_diagonal(A, diagonal);
}

// linear_operator
template <typename ValueType, typename MemorySpace>
template<typename MatrixType, typename VectorType1, typename VectorType2>
void multi_jacobi<ValueType,MemorySpace>
::operator()(const MatrixType& A, const VectorType1& b, VectorType2& x)
{
    multi_jacobi<ValueType,MemorySpace>::operator()(A,b,x,default_omega);
}

// override default omega
template <typename ValueType, typename MemorySpace>
template<typename MatrixType, typename VectorType1, typename VectorType2>
void multi_jacobi<ValueType,MemorySpace>
::operator()(const MatrixType& A, const VectorType1& b, VectorType2& x, ValueType omega)
{
    CUSP_PROFILE_SCOPED();

    int num_slices = A.getNumSlices();

    // y <- A*x
    /*#ifdef SYNC_MULT
    //ENABLE_P2P();
    OMP_PAR_FOR();
    for( int index = 0; index < num_slices; index++ )
    {
	cudaSetDevice(A.getDeviceNum(index));
        thrust::scatter(x(index).begin(), x(index).end(), A.local_rows[index].begin(), x2((index+1)%num_slices).begin());
        thrust::scatter(x(index).begin(), x(index).end(), A.local_rows[index].begin(), x2(index).begin());
    }
    //DISABLE_P2P();
    cusp::multiply(A,x2,temp);
    #else*/
    cusp::multiply(A,x,temp);
    //#endif

    // x <- x + omega * D^-1 * (b - y)

    for ( int index = 0; index < num_slices; index++ )
    {
	cudaSetDevice( A.getDeviceNum(index) );
        thrust::transform(thrust::make_zip_iterator(thrust::make_tuple( x(index).begin(), diagonal(index).begin(), b(index).begin(), temp(index).begin())),
                          thrust::make_zip_iterator(thrust::make_tuple(	x(index).end()  , diagonal(index).end()  , b(index).end()  , temp(index).end())),
                          x(index).begin(),
                          detail::jacobi_postsmooth_functor<ValueType>(omega));
    }
}

template <typename ValueType, typename MemorySpace>
template<typename MatrixType, typename VectorType1, typename VectorType2>
void multi_jacobi<ValueType,MemorySpace>
::presmooth(const MatrixType& A, const VectorType1& b, VectorType2& x)
{
    CUSP_PROFILE_SCOPED();

    // x <- omega * D^-1 * b
    int num_slices = A.getNumSlices();

    for ( int index = 0; index < num_slices; index++ )
    {
	cudaSetDevice( A.getDeviceNum(index) );
        thrust::transform(b(index).begin(), b(index).end(),
                          diagonal(index).begin(),
                          x(index).begin(),
                          detail::jacobi_presmooth_functor<ValueType>(default_omega));
   }
}

template <typename ValueType, typename MemorySpace>
template<typename MatrixType, typename VectorType1, typename VectorType2>
void multi_jacobi<ValueType,MemorySpace>
::postsmooth(const MatrixType& A, const VectorType1& b, VectorType2& x)
{
    CUSP_PROFILE_SCOPED();

    int num_slices = A.getNumSlices();

    /*#ifdef SYNC_MULT
    //ENABLE_P2P();
    OMP_PAR_FOR();
    for( int index = 0; index < num_slices; index++ )
    {
	cudaSetDevice(A.getDeviceNum(index));
        thrust::scatter(x(index).begin(), x(index).end(), A.local_rows[index].begin(), x2((index+1)%num_slices).begin());
        thrust::scatter(x(index).begin(), x(index).end(), A.local_rows[index].begin(), x2(index).begin());
    }
    //DISABLE_P2P();
    cusp::multiply(A,x2,temp);
    #else*/
    cusp::multiply(A,x,temp);
    //#endif

    // x <- x + omega * D^-1 * (b - y)

    for ( int index = 0; index < num_slices; index++ )
    {
	cudaSetDevice( A.getDeviceNum(index) );
        thrust::transform(thrust::make_zip_iterator(thrust::make_tuple(x(index).begin(), diagonal(index).begin(), b(index).begin(), temp(index).begin())),
                          thrust::make_zip_iterator(thrust::make_tuple(x(index).end(),   diagonal(index).end(),   b(index).end(),   temp(index).end())),
                          x(index).begin(),
                          detail::jacobi_postsmooth_functor<ValueType>(default_omega));
    }
}

} // end namespace relaxation
} // end namespace cusp

